execfile('v4/base.py')
from typing import Optional

class XlF(Base):
    def __init__(self):
        self.bar_count_default = 50
        self.end_date_str = '20250831'  # 如果需要可以取消注释
        # self.start_date_str = '2025-06-16'  # 如果需要可以取消注释

    def mark_xl(self,df, label_window=3, threshold=0.15):
        """
        生成预测三日后涨幅超15%概率的大单特征
        
        参数:
            df: 原始数据 (需含date,code,close,xl,xl_rate等字段)
            label_window: 预测窗口天数 (默认3日)
            threshold: 涨幅阈值 (默认15%)
        
        返回:
            含特征和标签的DataFrame
        """
        df = df.sort_values(['code', 'date']).copy()
        
        # ===== 标签生成 =====
        df['future_close'] = df.groupby('code')['close'].shift(-label_window)
        
        # ===== 大单特征工程 =====
        # 1. 基础特征
        df['xl_net'] = df['xl']
        df['xl_ratio'] = df['xl_rate']
    
        def _cal_yesterday_feature(df):
            df = df.sort_values(by=['code', 'date'])
            df['pre_xl_ratio'] = df.groupby('code')['xl_ratio'].shift(1)
            return df
        df = df.groupby('code', group_keys=False)[df.columns].apply(_cal_yesterday_feature)
        
        
        # 2. 时间窗口特征
        windows = [3,5,10]
        for w in windows:
            # 滚动窗口计算
            df[f'xl_cumsum_{w}d'] = df.groupby('code')['xl'].transform(lambda s: s.rolling(w).sum())
            df[f'xl_rate_ma_{w}d'] = df.groupby('code')['xl_rate'].transform(lambda s: s.rolling(w).mean())
            df[f'xl_vol_{w}d'] = df.groupby('code')['xl'].transform(lambda s: s.rolling(w).std())
            
        # 3. 量价协同特征
        df['xl_push_eff'] = df['xl'] / (df['turnover'] * df['quote_rate'].abs().replace(0, 0.001))
        df['xl_price_div'] = (df['xl'] > df['xl'].rolling(5).mean()) & (df['close'] < df['close'].rolling(5).mean())
        
        # 4. 异常波动特征
        df['xl_spike'] = (df['xl'] > 2 * df.groupby('code')['xl'].transform(lambda s: s.rolling(20).mean()))
        df['xl_break_high'] = (df['xl'] > df.groupby('code')['xl'].transform(lambda s: s.rolling(10).max()))
        
        # 5. 资金流特征
        df['xl_concentration'] = df['xl'] / df.groupby('code')['xl'].transform('mean')
        df['xl_t_rate'] = df['xl'] / (df['turnover'] / df['close'])
        
        # 6. 持续性特征
        df['xl_positive_days'] = df.groupby('code')['xl'].transform(
            lambda s: s.gt(0).groupby(s.diff().ne(0).cumsum()).cumcount() + 1
        )
        
        # ===== 清理中间列 =====
        df = df.drop(columns=['future_close'])
        df = self.drop_base_columns(df)
        return df

    def mark_xl_col_growth(
        self,
        df: pd.DataFrame,
        target_col: str,
        window: int = 10,
        feature_suffix: Optional[str] = None,
        min_valid_value: float = 1e-6,
        growth_type: str = 'relative',
        trading_col: str = 'is_trading',
        is_drop: bool = True
    ) -> pd.DataFrame:
        """
        优化版增长率计算 - 保持原有4列输出
        
        改进点：
        1. 零值特殊处理：当基准值为零时，使用替代策略计算增长率
        2. 仅保留最终需要的4列输出
        3. 使用临时DataFrame避免污染原始数据
        """
        # 输入校验
        required_cols = {'code', 'date', target_col}
        missing_cols = required_cols - set(df.columns)
        if missing_cols:
            raise ValueError(f"缺失必要列: {missing_cols}")
        
        # 创建临时工作DataFrame（仅包含必要列）
        work_df = df[['code', 'date', target_col]].copy()
        
        # 数据预处理
        work_df['date'] = pd.to_datetime(work_df['date'])
        work_df.sort_values(by=['code', 'date'], inplace=True)
        
        # 特征列命名
        suffix = f"_{feature_suffix}" if feature_suffix else ''
        growth_col = f"{target_col}_yoy{suffix}_{window}d"
        pre_growth_col = f"pre_{growth_col}"
        
        # 零值特殊处理 - 添加微小常数避免除零错误
        epsilon = 1e-10
        work_df['_adjusted_target'] = work_df[target_col].replace(0, epsilon) + epsilon
        
        # 滚动基准计算（优化零值处理）
        work_df['_base'] = work_df.groupby('code', group_keys=False)['_adjusted_target'].transform(
            lambda x: x.shift(1).rolling(window, min_periods=1).mean()
        )
        # display(work_df[-30:])
        # 增长率计算（优化零值处理）
        prev_col = work_df.groupby('code')[target_col].shift(1)
        
        if growth_type == 'relative':
            # 核心优化：零值特殊处理
            with np.errstate(divide='ignore', invalid='ignore'):
                work_df[growth_col] = np.where(
                    work_df['_base'] <= min_valid_value,
                    # 基准接近零时的替代计算
                    (work_df[target_col] - prev_col) / (prev_col + epsilon),
                    # 正常相对增长率计算
                    (work_df[target_col] - work_df['_base']) / work_df['_base']
                )
            # display(work_df[-30:])
        elif growth_type == 'log':
            # 对数增长率同样需要零值处理
            work_df[growth_col] = np.where(
                (work_df['_base'] > min_valid_value) & (work_df[target_col] > min_valid_value),
                np.log(work_df[target_col] / work_df['_base']),
                np.where(
                    # 零到非零的跃迁
                    (work_df[target_col] > min_valid_value) & (prev_col <= min_valid_value),
                    1.0,  # 固定值表示显著增长
                    np.nan
                )
            )
        else:
            raise ValueError("growth_type 必须为 'relative' 或 'log'")
        
        # 计算昨日特征
        work_df[pre_growth_col] = work_df.groupby('code')[growth_col].shift(1)
        
        # 仅保留最终需要的4列
        result = work_df[['code', 'date', growth_col, pre_growth_col]].copy()
        
        # 停牌数据处理（如果提供停牌列）
        if trading_col in df.columns:
            # 合并停牌信息
            trading_info = df[['code', 'date', trading_col]].drop_duplicates()
            result = result.merge(trading_info, on=['code', 'date'], how='left')
            result.loc[~result[trading_col], [growth_col, pre_growth_col]] = np.nan
            result.drop(columns=[trading_col], inplace=True)
        
        return result.round(4)  # 提高精度保留小数